// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                            avtCubeFileFormat.C                           //
// ************************************************************************* //

#include <avtCubeFileFormat.h>

#include <string>

#include <vtkFloatArray.h>
#include <vtkRectilinearGrid.h>
#include <vtkStructuredGrid.h>
#include <vtkUnstructuredGrid.h>
#include <vtkPolyData.h>

#include <avtDatabaseMetaData.h>

#include <DBOptionsAttributes.h>
#include <Expression.h>

#include <InvalidVariableException.h>
#include <InvalidDBTypeException.h>
#include <DebugStream.h>
#include <iostream>
#include <fstream>

using std::string;
using std::stringstream;

// ****************************************************************************
//  Method: avtCubeFileFormat constructor
//
//  Programmer: prabhat -- generated by xml2avt
//  Creation:   Fri Jun 12 09:11:12 PDT 2009
//
//  Programmer: oruebel -- Rename to Cube and removed Walker functionality
//  Modified:   Thu May 12 11:18 PDT 2009
//
//  Programmer: jfavre -- read boolean option
//  Modified:   Fri Apr 26 11:12:38 CEST 2013
// ****************************************************************************

avtCubeFileFormat::avtCubeFileFormat(const char *filename, const DBOptionsAttributes *readOpts)
    : avtMTSDFileFormat(&filename, 1)
{
  ExtendVolumeByOneCell = readOpts->GetBool("ExtendVolumeByOneCell");
  creader = new CubeReader(filename);
  if( !creader->isFileValid() ){
    string ex = "File is not a valid CUBE file. ";
    ex += creader->getErrorString();
    EXCEPTION1(InvalidDBTypeException,  ex.c_str() );
  }
}

// ****************************************************************************
//  Method: avtEMSTDFileFormat::GetNTimesteps
//
//  Purpose:
//      Tells the rest of the code how many timesteps there are in this file.
//
//  Programmer: prabhat -- generated by xml2avt
//  Creation:   Fri Jun 12 09:11:12 PDT 2009
//
// ****************************************************************************

int
avtCubeFileFormat::GetNTimesteps(void)
{
    return creader->GetNTimesteps();
}


// ****************************************************************************
//  Method: avtCubeFileFormat::FreeUpResources
//
//  Purpose:
//      When VisIt is done focusing on a particular timestep, it asks that
//      timestep to free up any resources (memory, file descriptors) that
//      it has associated with it.  This method is the mechanism for doing
//      that.
//
//  Programmer: prabhat -- generated by xml2avt
//  Creation:   Fri Jun 12 09:11:12 PDT 2009
//
// ****************************************************************************

void
avtCubeFileFormat::FreeUpResources(void)
{
}


// ****************************************************************************
//  Method: avtCubeFileFormat::PopulateDatabaseMetaData
//
//  Purpose:
//      This database meta-data object is like a table of contents for the
//      file.  By populating it, you are telling the rest of VisIt what
//      information it can request from you.
//
//  Programmer: prabhat -- generated by xml2avt
//  Creation:   Fri Jun 12 09:11:12 PDT 2009
//
//  Modifications:
//    Jean Favre, Fri Feb  1 11:30:09 CET 2013
//    added support to set the Unit Cell Vector
//
//  Programmer: jfavre -- use the new boolean option to get the Unit Cell
//  Modified:   Fri Apr 26 11:12:38 CEST 2013
// ****************************************************************************

void
avtCubeFileFormat::PopulateDatabaseMetaData(avtDatabaseMetaData *md, int timeState)
{
  if (creader) {
    
    string meshname = "electron_grid";
    int nblocks = 1;  
    int block_origin = 0;
    int spatial_dimension = 3;
    int topological_dimension = 3;
    //double *extents = NULL;
    avtMeshType mt;

    bool shear = creader->isGridSheared();
    debug4 << "getMesh() shear= " << shear << endl;
    
    if (shear) 
      mt = AVT_CURVILINEAR_MESH;
    else 
      mt = AVT_RECTILINEAR_MESH;
    
    avtMeshMetaData *mesh = new avtMeshMetaData(meshname, nblocks, block_origin, 0, 0,
                                               spatial_dimension, topological_dimension,
                                               mt);
    mesh->hasSpatialExtents = false;
    creader->GetUnitCell(mesh->unitCellOrigin, mesh->unitCellVectors, ExtendVolumeByOneCell);
    md->Add(mesh);
    //AddMeshToMetaData(md, meshname, mt, extents, nblocks, block_origin, 
                      //spatial_dimension, topological_dimension);

    string mesh_for_this_var = "electron_grid";
    string varname = "electron_density";
    avtCentering cent = AVT_NODECENT;
    
    int num_orbitals = creader->GetNumOrbitals();
    if (num_orbitals==0) {
      varname = "electron_density"; 
      AddScalarVarToMetaData(md, varname, mesh_for_this_var, cent);
    }
    else {
      for (int i=0; i<num_orbitals; i++) {
        int orbit = creader->GetOrbitalNumber(i);
        stringstream ss;
        ss<<"orbital_"<<orbit;
        varname = ss.str();
        AddScalarVarToMetaData(md, varname, mesh_for_this_var, cent);
      }

    }


    meshname = "atoms";
    mt = AVT_POINT_MESH;
    
    nblocks = 1;  
    block_origin = 0;
    spatial_dimension = 3;
    topological_dimension = 0;
    //extents = NULL;
    
    avtMeshMetaData *mesh2 = new avtMeshMetaData(meshname, nblocks, block_origin, 0, 0,
                                               spatial_dimension, topological_dimension,
                                               mt);
    mesh2->hasSpatialExtents = false;
    creader->GetUnitCell(mesh2->unitCellOrigin, mesh2->unitCellVectors, ExtendVolumeByOneCell);
    md->Add(mesh2);
    //AddMeshToMetaData(md, meshname, mt, extents, nblocks, block_origin,
                      //spatial_dimension, topological_dimension);
    
    
    mesh_for_this_var = "atoms";
    varname = "element";
    cent = AVT_NODECENT;
    AddScalarVarToMetaData(md, varname, mesh_for_this_var, cent);
    
  }
}


// ****************************************************************************
//  Method: avtCubeFileFormat::GetMesh
//
//  Purpose:
//      Gets the mesh associated with this file.  The mesh is returned as a
//      derived type of vtkDataSet (ie vtkRectilinearGrid, vtkStructuredGrid,
//      vtkUnstructuredGrid, etc).
//
//  Arguments:
//      timestate   The index of the timestate.  If GetNTimesteps returned
//                  'N' time steps, this is guaranteed to be between 0 and N-1.
//      meshname    The name of the mesh of interest.  This can be ignored if
//                  there is only one mesh.
//
//  Programmer: prabhat -- generated by xml2avt
//  Creation:   Fri Jun 12 09:11:12 PDT 2009
//
//  Programmer: jfavre -- use GetShearedGridLocations2()
//  Modified:   Fri Apr 26 11:12:38 CEST 2013
// ****************************************************************************

vtkDataSet *
avtCubeFileFormat::GetMesh(int timestate, const char *meshname)
{
  int nnodes;
  float *xarray = NULL, *yarray = NULL, *zarray = NULL; ///TODO: check on fix for uninitialized warning
  if (creader) {

    if (strcmp(meshname,"electron_grid")==0) {
      
      bool shear = creader->isGridSheared();
      debug4 << "sheared = " << shear << endl;

      if (shear) {
        //int ndims = 3;
        int dims[3] = {1,1,1};

        debug4 << "Creating curvilinear grid" << endl;
        creader->GetGridSize(dims[0],dims[1],dims[2]);
        if(ExtendVolumeByOneCell)
        {
          dims[0]++;
          dims[1]++;
          dims[2]++;
          nnodes = dims[0]*dims[1]*dims[2];
          // no coordinate array allocation because we will allocate the point array directly and use it
        }
        else
        {
          nnodes = dims[0]*dims[1]*dims[2];

          debug4 << "\tgrid dims are " << dims[0] << ", " << dims[1] << ", " << dims[2] << endl;

          xarray = new float[nnodes];
          yarray = new float[nnodes];
          zarray = new float[nnodes];

          creader->GetShearedGridLocations(xarray, yarray, zarray);
          debug4 << "\tdone reading sheared grid locations" << endl;
        }

        debug4 << "\ttrying to allocate grid" << endl;
        vtkStructuredGrid *sgrid = vtkStructuredGrid::New();
        vtkPoints *points = vtkPoints::New();
        sgrid->SetPoints(points);
        sgrid->SetDimensions(dims);
        points->Delete();
        points->SetNumberOfPoints(nnodes);
        debug4 << "\tdone allocating, now filling in values" << endl;

        float *pts = (float *) points->GetVoidPointer(0);
        if(ExtendVolumeByOneCell)
        {
          creader->GetShearedGridLocations2(pts);
        }
        else
        {
          float *xc = xarray;
          float *yc = yarray;
          float *zc = zarray;

          for(int k = 0; k < dims[2]; ++k)
            for(int j = 0; j < dims[1]; ++j)
              for(int i = 0; i < dims[0]; ++i)
              {
                *pts++ = *xc++;
                *pts++ = *yc++;
                *pts++ = *zc++;
              }
          debug4 << "\tdone with setting pts" << endl;

          delete []xarray;
          delete []yarray;
          delete []zarray;
        }


        debug4 << "returning sgrid" << endl;
        return sgrid;
      }//end sheared mesh
      else //mesh not sheared 
      { 
        debug4 << "Creating a rectilinear mesh" << endl;
        int xs, ys, zs;
        creader->GetGridSize(xs,ys,zs);

        vtkFloatArray *coords[3] = {0,0,0};

        coords[0] = vtkFloatArray::New();
        coords[0]->SetNumberOfTuples(xs);
        float *xarray = (float*)coords[0]->GetVoidPointer(0);

        coords[1] = vtkFloatArray::New();
        coords[1]->SetNumberOfTuples(ys);
        float *yarray = (float*)coords[1]->GetVoidPointer(0);

        coords[2] = vtkFloatArray::New();
        coords[2]->SetNumberOfTuples(zs);
        float *zarray = (float*)coords[2]->GetVoidPointer(0);

        creader->GetGridLocations(xarray, yarray, zarray);

        vtkRectilinearGrid *rgrid = vtkRectilinearGrid::New();

        rgrid->SetDimensions(xs, ys, zs);

        rgrid->SetXCoordinates(coords[0]);
        rgrid->SetYCoordinates(coords[1]);
        rgrid->SetZCoordinates(coords[2]);

        coords[0]->Delete();
        coords[1]->Delete();
        coords[2]->Delete();    

        return rgrid;
      }//end unsheared mesh
    } //end read electron grid
    else if (strcmp(meshname,"atoms")==0) {
      
      //debug4 << "Creating static atoms..." << endl;
      
      vtkPolyData *dataset = vtkPolyData::New();
      vtkPoints *vtkpoints = vtkPoints::New();
      
      int natoms = creader->GetNAtoms();
      std::vector<float> atom_locations;
      int num_dims = 3;
      atom_locations.resize(natoms*num_dims);
    
      creader->GetAtomLocations(atom_locations); 
      
      vtkpoints->SetNumberOfPoints((vtkIdType) natoms);
      
      float *pts = (float *) vtkpoints->GetVoidPointer(0);
      
      for (int i=0; i < natoms*num_dims; i++)
        pts[i] = atom_locations[i];
      
      dataset->Allocate(natoms*num_dims);
      for (long int i=0; i < natoms; i++){
        vtkIdType onevertex = (vtkIdType) i;
        dataset->InsertNextCell(VTK_VERTEX, 1, &onevertex);
      }
      
      dataset->SetPoints(vtkpoints);
      vtkpoints->Delete();
      
      return dataset;
    }//read atoms

    
  }//cubereader

  return NULL;
}


// ****************************************************************************
//  Method: avtCubeFileFormat::GetVar
//
//  Purpose:
//      Gets a scalar variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      timestate  The index of the timestate.  If GetNTimesteps returned
//                 'N' time steps, this is guaranteed to be between 0 and N-1.
//      varname    The name of the variable requested.
//
//  Programmer: prabhat -- generated by xml2avt
//  Creation:   Fri Jun 12 09:11:12 PDT 2009
//
//  Programmer: jfavre -- use GetGridValues2()
//  Modified:   Fri Apr 26 11:12:38 CEST 2013
// ****************************************************************************

vtkDataArray *
avtCubeFileFormat::GetVar(int timestate, const char *varname)
{
  //debug4 << "Trying to get data for "<<varname<<endl;
  
  if (creader)
    {
    debug4<< "GetVar() called for " << varname << endl;

    if (strcmp(varname, "electron_density")==0)
      {
      int ntuples, x, y, z;
      creader->GetGridSize(x, y, z);
      if(ExtendVolumeByOneCell)   
        ntuples = (++x)*(++y)*(++z);
     else
        ntuples = x*y*z;
      vtkFloatArray *rv = vtkFloatArray::New();
      rv->SetNumberOfTuples(ntuples);
      float* ptr = (float*)rv->GetVoidPointer(0);
      if(ExtendVolumeByOneCell)
      {
        //for (int all=0; all < x*y*z; all++)
          //ptr[all] = 0;
        creader->GetGridValues2(ptr);
      }
      else
      {
        creader->GetGridValues(ptr);
      }
      return rv;
      }
    else if (strstr(varname,"orbital")!=NULL)
      {
      int x,y,z;
      creader->GetGridSize(x,y,z);
      
      int ntuples = x*y*z;
      debug4 << "\t x=" << x << ", y=" << y << ", z=" << z << endl;
      
      vtkFloatArray *rv = vtkFloatArray::New();
      rv->SetNumberOfTuples(ntuples);
      float* ptr = (float*)rv->GetVoidPointer(0);
      
      creader->GetOrbitalValues(ptr, varname);
      return rv;
      }
    else if (strcmp(varname, "element")==0)
      {
      int natoms = creader->GetNAtoms();
      int ntuples = natoms;
      vtkFloatArray *rv = vtkFloatArray::New();
      rv->SetNumberOfTuples(ntuples);
      float* ptr = (float*)rv->GetVoidPointer(0);
      
      creader->GetAtomTypes(ptr);
      return rv;

      }
    else
      {
      debug4 << "cubereader WARNING:: want data for "<< varname<<endl;
      }
  }
  
  return NULL;
}


// ****************************************************************************
//  Method: avtCubeFileFormat::GetVectorVar
//
//  Purpose:
//      Gets a vector variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      timestate  The index of the timestate.  If GetNTimesteps returned
//                 'N' time steps, this is guaranteed to be between 0 and N-1.
//      varname    The name of the variable requested.
//
//  Programmer: prabhat -- generated by xml2avt
//  Creation:   Fri Jun 12 09:11:12 PDT 2009
//
// ****************************************************************************

vtkDataArray *
avtCubeFileFormat::GetVectorVar(int timestate, const char *varname)
{
    return NULL;
}
